import json
import os
import re
from pysmi.codegen import JsonCodeGen
from pysmi.compiler import MibCompiler
from pysmi.parser import SmiStarParser
from pysmi.reader import FileReader
from pysmi.searcher import StubSearcher
from pysmi.writer import CallbackWriter


'''
编写时间 2024 08 08 

用来将mib文件，转换成json

pysmi 1.4.4
Python 3.10.14 

'''

def dirTraversal(dir):
    '''
    文件递归遍历生成器
    :param dir:
    :return:
    '''
    import os
    if not os.path.isdir(dir):
        raise RuntimeError(dir + ' is not a dir.')
    try:
        for item in os.listdir(dir):
            fin_item = {}
            if os.path.isdir(dir + os.path.sep + item):
                fin_item['type'] = 'dir'
                fin_item['path'] = dir + os.path.sep + item
                yield fin_item
                yield from dirTraversal(fin_item['path'])
            else:
                fin_item['type'] = 'file'
                fin_item['path'] = dir + os.path.sep + item
                yield fin_item
    except PermissionError as e:
        print(e.__str__())


def getModuleName(filePath):
    '''
    读取mib文件内容，找到定义部分。
    获取模块的定义部分，取模块的名字。
    :param filePath:  文件绝对路径
    :return:  如果没找到模块名，就返回none
    '''
    file = open(filePath, 'r')
    pattern_src = r'.*DEFINITIONS.*BEGIN'
    pattern = re.compile(pattern_src)
    for line in file:
        result = pattern.match(line)
        if result != None:
            return result[0].strip().split(" ")[0]

def getMibFileList(srcDir):
    '''
    获取后缀名为 mib的文件列表，绝对路径
    :param srcDir:
    :return:
    '''

    result = []
    for file in dirTraversal(srcDir):
        if file["type"] == "file":
            # 后缀名
            postFixName = file["path"].split(os.sep)[-1].split(".")[-1]
            filePath = file["path"]
            # 过滤后缀名
            if postFixName == "mib":
                result.append(filePath)

    if len(result) != 0:
        return result
    else:
        return None

def getModuleNameList(MibFileList):
    '''
    获取指定文件列表中的mib文件定义名列表
    :param MibFileList:  mib文件列表
    :return:
    '''

    moduleNameList = []
    for filePath in MibFileList:
        moduleName = getModuleName(filePath)
        if moduleName != None:
            moduleNameList.append(moduleName)
        else:
            print("Error: getModuleName 失败---->",filePath)
    return moduleNameList

def renameMibFiles(srcDir):
    '''
    递归目标目录，将文件名和模块名对不上的文件，重新以模块名命名mib文件。
    :param srcDir:
    :return:
    '''

    # 要处理的目录，里边可以任意组织结构，这里是递归遍历后缀名为mib的文件。
    # srcDir = "E:\DOC\MibFiles"

    for file in dirTraversal(srcDir):
        if file["type"] == "file":
            # 后缀名
            postFixName = file["path"].split(os.sep)[-1].split(".")[-1]
            # 无后缀的文件名
            fileBaseName = file["path"].split(os.sep)[-1].split(".")[0]
            # 文件名
            fileName = file["path"].split(os.sep)[-1]
            # 文件绝对路径
            filePath = file["path"]
            fileDir = os.sep.join(file["path"].split(os.sep)[:-1])
            # 过滤后缀名
            if postFixName == "mib":
                # 获取模块名
                moduleName = getModuleName(filePath)
                if moduleName != None:
                    # 如果文件名不规范，则重命名
                    if fileBaseName != moduleName:
                        print("------------------------------------------------------------")
                        print(filePath, "模块名和文件名不一致。")
                        newFilePath = ".".join([os.sep.join([fileDir, moduleName]), "mib"])
                        if not os.path.exists(newFilePath):
                            print("新文件路径可用。尝试重命名mib文件。")
                            print("newFilePath:", newFilePath)
                            os.rename(filePath, newFilePath)
                        else:
                            print(newFilePath, "已经存在。")
                            print("！！！ERROR:   新文件路径不可用 ！！！")
                else:
                    print("Error: getModuleName 失败---->", filePath)

    print("执行完成，如果无输出内容，说明文件命名全部规范。")

def getComplyFailList(srcdDir):

    '''
    尝试直接编译mib目录中的库，看看有多少文件编译失败。
    :param srcdDir:
    :return:
    '''
    print(f"尝试对目录{srcdDir}进行mib文件编译，时间可能较久。请耐心等待。")

    failList = []
    mibFileList = getMibFileList(srcdDir)
    moduleList = getModuleNameList(mibFileList)

    # 定义要编译的 MIB 文件名称列表
    inputMibs = moduleList

    # 定义一个回调函数，用于打印编译后的 JSON 文档
    def printOut(mibName, jsonDoc, cbCtx):
        pass

    # 初始化编译器基础设施
    mibCompiler = MibCompiler(SmiStarParser(), JsonCodeGen(), CallbackWriter(printOut))

    mibCompiler.addSources(FileReader(srcDir))

    # 避免重新编译包含宏的 MIB 文件
    mibCompiler.addSearchers(StubSearcher(*JsonCodeGen.baseMibs))

    # 执行递归的 MIB 编译过程
    results = mibCompiler.compile(*inputMibs)

    for k in results:
        if results[k] == "failed":
            failList.append(k)

    return failList



def complyOneMib(moduleName,dependDir,destDir):
    '''
    编译一个mib文件
    :param moduleName: 文件内容中的定义名，也就是模块名
    :param dependDir:  存放mib文件依赖的目录，通常是一个mib文件包。pysmi会递归依赖
    :param destDir:  输出json文件的目录。
    :return:  none
    '''


    if not os.path.exists(destDir):
        raise RuntimeError("complyOneMib 目标目录不存在",destDir)
    elif not os.path.isdir(destDir):
        raise RuntimeError("complyOneMib 目标目录并非目录",destDir)

    # 定义要编译的 MIB 文件名称列表
    inputMibs = [moduleName]

    # 定义一个回调函数，用于打印编译后的 JSON 文档
    def printOut(mibName, jsonDoc, cbCtx):
        jsonFilePath = destDir+os.sep+mibName+".json"
        jsonFile = open(jsonFilePath,'w')
        jsonFile.write(jsonDoc)
        jsonFile.flush()
        jsonFile.close()
    # 初始化编译器基础设施
    mibCompiler = MibCompiler(SmiStarParser(), JsonCodeGen(), CallbackWriter(printOut))

    mibCompiler.addSources(FileReader(dependDir))

    # 避免重新编译包含宏的 MIB 文件
    mibCompiler.addSearchers(StubSearcher(*JsonCodeGen.baseMibs))

    # 执行递归的 MIB 编译过程
    mibCompiler.compile(*inputMibs)


def complyMibs(srcDir,destDir,dependDir=None):

    '''
    将指定的目录中的mib文件，编译成json文件。
    1.文件名后缀需要是.mib
    2.如果不是，则先运行renameMibFiles 方法重命名mib文件
    :param srcDir:  要编译的mib文件目录
    :param destDir:  输出的json文件的目录
    :param dependDir:  mib依赖文件目录，可以和srcDir不同，如果不指定则默认相同
    :return:  无
    '''

    #如果依赖mib库的路径没有设置，那么就用源mib路径作为依赖库
    if dependDir == None:
        dependDir = srcDir

    fileList = getMibFileList(srcDir)
    moduleNameList = getModuleNameList(fileList)

    for name in moduleNameList:
        complyOneMib(name, dependDir, destDir)



if __name__ == '__main__':
    # srcDir = "E:\DOC\MibFiles"
    srcDir = r'E:\DOC\MibFiles\H3c\H3C New Style Private MIB'
    destDir = "E:\Desk\mibtest"
    dependDir = "E:\DOC\MibFiles"

    complyMibs(srcDir,destDir,dependDir)